{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "<a href=\"https://colab.research.google.com/drive/1Qa9KWk0upKcLBqCwP66z6nzLP4TDnXE-?usp=sharing\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"></a>"
      ],
      "metadata": {
        "id": "U_LAtyCVZYjS"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### ART: Automatic Reasoning and Tool-use Agent\n",
        "\n",
        "An agent that automatically selects and executes tools through:\n",
        "1. Parse task library for similar tasks\n",
        "2. Select relevant tool demonstrations\n",
        "3. Generate reasoning steps\n",
        "4. Execute actual tool calls at appropriate steps\n",
        "5. Adapt based on tool execution results"
      ],
      "metadata": {
        "id": "bpoJNyzlZZ4g"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "ToJymutGZXLv"
      },
      "outputs": [],
      "source": [
        "!pip install -qU google-generativeai"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import google.generativeai as genai\n",
        "import getpass\n",
        "import re\n",
        "from datetime import datetime"
      ],
      "metadata": {
        "id": "aazlKXFLZtnN"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Get free-tier Google's Gemini API Key here: https://aistudio.google.com/app/apikey"
      ],
      "metadata": {
        "id": "bLG0g7evZyaj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "API_KEY = getpass.getpass(\"Enter your Google API key: \")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RI1H_5ggZwCx",
        "outputId": "06ae1c6e-3dd0-4223-ec4e-fb84a19122c6"
      },
      "execution_count": 3,
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Enter your Google API key: ··········\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "genai.configure(api_key=API_KEY)"
      ],
      "metadata": {
        "id": "hUCiZFW9Z7Aw"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class ARTAgent:\n",
        "    def __init__(self):\n",
        "        self.model = genai.GenerativeModel(\"gemini-2.0-flash\")\n",
        "        self.tools = {}\n",
        "        self.task_library = []\n",
        "\n",
        "    def add_tool(self, name, func, description, param_name):\n",
        "        self.tools[name] = {\"func\": func, \"desc\": description, \"param\": param_name}\n",
        "\n",
        "    def add_example(self, task, tools_used):\n",
        "        \"\"\"Add successful task example\"\"\"\n",
        "        self.task_library.append({\"task\": task, \"tools\": tools_used})\n",
        "\n",
        "    def find_similar(self, query):\n",
        "        \"\"\"Find similar past tasks\"\"\"\n",
        "        if not self.task_library:\n",
        "            return []\n",
        "\n",
        "        # Simple keyword matching\n",
        "        words = set(query.lower().split())\n",
        "        similar = []\n",
        "        for ex in self.task_library:\n",
        "            ex_words = set(ex[\"task\"].lower().split())\n",
        "            if words & ex_words:  # Any overlap\n",
        "                similar.append(ex)\n",
        "        return similar[:2]\n",
        "\n",
        "    def plan(self, query, similar):\n",
        "        \"\"\"Generate step-by-step plan\"\"\"\n",
        "        tools_list = \"\\n\".join([\n",
        "            f\"{n}({t['param']}): {t['desc']}\"\n",
        "            for n, t in self.tools.items()\n",
        "        ])\n",
        "\n",
        "        examples = \"\\n\".join([\n",
        "            f\"'{ex['task']}' used: {', '.join(ex['tools'])}\"\n",
        "            for ex in similar\n",
        "        ]) if similar else \"No examples.\"\n",
        "\n",
        "        prompt = f\"\"\"Create a step-by-step plan for this task.\n",
        "\n",
        "        Task: {query}\n",
        "\n",
        "        Tools (use exact parameter names):\n",
        "        {tools_list}\n",
        "\n",
        "        Similar tasks:\n",
        "        {examples}\n",
        "\n",
        "        Format each step as:\n",
        "        Step N: description [TOOL: name(param=\"value\")] or [TOOL: None]\n",
        "        Use ONLY the parameter names shown above.\n",
        "\n",
        "        Plan:\"\"\"\n",
        "\n",
        "        return self.model.generate_content(prompt).text\n",
        "\n",
        "    def execute(self, query):\n",
        "        \"\"\"Execute task with automatic tool selection\"\"\"\n",
        "        print(f\"\\n🎯 Task: {query}\\n\")\n",
        "\n",
        "        # Find similar tasks\n",
        "        similar = self.find_similar(query)\n",
        "        if similar:\n",
        "            print(f\"📚 Found {len(similar)} similar tasks\\n\")\n",
        "\n",
        "        # Generate plan\n",
        "        plan_text = self.plan(query, similar)\n",
        "        print(f\"📋 Plan:\\n{plan_text}\\n\")\n",
        "\n",
        "        # Execute plan\n",
        "        tools_used = []\n",
        "        results = []\n",
        "\n",
        "        for line in plan_text.split(\"\\n\"):\n",
        "            if not line.strip().startswith(\"Step\"):\n",
        "                continue\n",
        "\n",
        "            print(f\"▶ {line.split(':')[0]}\")\n",
        "\n",
        "            # Check for tool call\n",
        "            tool_match = re.search(r'\\[TOOL:\\s*(\\w+)\\((.*?)\\)\\]', line)\n",
        "            if tool_match:\n",
        "                tool_name = tool_match.group(1)\n",
        "\n",
        "                if tool_name == \"None\":\n",
        "                    print(\"  (reasoning only)\\n\")\n",
        "                    continue\n",
        "\n",
        "                # Parse params - take first param only\n",
        "                param_name = self.tools[tool_name][\"param\"]\n",
        "                param_value = None\n",
        "\n",
        "                param_match = re.search(r'\\w+=[\"\\'](.*?)[\"\\']', tool_match.group(2))\n",
        "                if param_match:\n",
        "                    param_value = param_match.group(1)\n",
        "\n",
        "                # Execute with correct param name\n",
        "                try:\n",
        "                    if param_value:\n",
        "                        result = self.tools[tool_name][\"func\"](**{param_name: param_value})\n",
        "                    else:\n",
        "                        result = self.tools[tool_name][\"func\"]()\n",
        "                    print(f\"  ✅ {tool_name}: {result}\\n\")\n",
        "                    results.append(result)\n",
        "                    tools_used.append(tool_name)\n",
        "                except Exception as e:\n",
        "                    print(f\"  ❌ Error: {e}\\n\")\n",
        "\n",
        "        # Learn from success\n",
        "        if results:\n",
        "            self.add_example(query, tools_used)\n",
        "            print(f\"✅ Completed. Used: {', '.join(tools_used)}\\n\")\n",
        "\n",
        "        return results"
      ],
      "metadata": {
        "id": "Llxzj7rKaFn_"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Define simple tools\n",
        "def fetch(source):\n",
        "    data = {\n",
        "        \"sales\": {\"revenue\": 150000, \"units\": 1200},\n",
        "        \"customers\": {\"total\": 5000, \"active\": 3500},\n",
        "    }\n",
        "    return data.get(source.lower(), {})\n",
        "\n",
        "def calc(expr):\n",
        "    return eval(expr)\n",
        "\n",
        "def analyze(data):\n",
        "    return f\"Analyzed: {len(data)} metrics\"\n",
        "\n",
        "def report(msg):\n",
        "    return f\"Report: {msg}\""
      ],
      "metadata": {
        "id": "hnKxzyFaaKlb"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Usage\n",
        "agent = ARTAgent()\n",
        "\n",
        "agent.add_tool(\"fetch\", fetch, \"Get data from source\", \"source\")\n",
        "agent.add_tool(\"calc\", calc, \"Calculate expression\", \"expr\")\n",
        "agent.add_tool(\"analyze\", analyze, \"Analyze data\", \"data\")\n",
        "agent.add_tool(\"report\", report, \"Create report\", \"msg\")\n",
        "\n",
        "# Add examples\n",
        "agent.add_example(\"analyze sales\", [\"fetch\", \"analyze\", \"report\"])\n",
        "\n",
        "# Run\n",
        "agent.execute(\"Get sales data and calculate total revenue\")\n",
        "agent.execute(\"Analyze customer data\")\n",
        "agent.execute(\"Get sales and create report\")  # Should use learned pattern\n",
        "\n",
        "print(f\"📚 Learned {len(agent.task_library)} tasks\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "jRj1FZYXaNB6",
        "outputId": "fe7f9d69-4875-4edb-f3af-51b54f0c3661"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "🎯 Task: Get sales data and calculate total revenue\n",
            "\n",
            "📚 Found 1 similar tasks\n",
            "\n",
            "📋 Plan:\n",
            "Okay, here's a step-by-step plan to get sales data and calculate total revenue, following the provided format and constraints:\n",
            "\n",
            "**Plan:**\n",
            "\n",
            "**Step 1**: Fetch sales data from the designated source. [TOOL: fetch(source=\"sales_data\")]\n",
            "**Step 2**: Calculate the total revenue using the fetched sales data.  Assume the 'sales_data' contains a field named 'price' and a field named 'quantity'. The expression calculates the sum of (price * quantity). [TOOL: calc(expr=\"SUM(price * quantity)\")]\n",
            "**Step 3**: Generate a report indicating the total revenue calculated. [TOOL: report(msg=\"Total revenue calculated\")]\n",
            "\n",
            "\n",
            "\n",
            "🎯 Task: Analyze customer data\n",
            "\n",
            "📚 Found 1 similar tasks\n",
            "\n",
            "📋 Plan:\n",
            "Okay, here's a step-by-step plan to analyze customer data, using the specified tools and their parameter names:\n",
            "\n",
            "**Plan:**\n",
            "\n",
            "Step 1: Retrieve the customer data from the specified source. [TOOL: fetch(source=\"customer_database\")]\n",
            "Step 2: Perform the analysis on the retrieved customer data. [TOOL: analyze(data=\"fetched_data\")]\n",
            "Step 3: Generate a report summarizing the key findings from the customer data analysis. [TOOL: report(msg=\"analysis_results\")]\n",
            "\n",
            "\n",
            "▶ Step 1\n",
            "  ✅ fetch: {}\n",
            "\n",
            "▶ Step 2\n",
            "  ✅ analyze: Analyzed: 12 metrics\n",
            "\n",
            "▶ Step 3\n",
            "  ✅ report: Report: analysis_results\n",
            "\n",
            "✅ Completed. Used: fetch, analyze, report\n",
            "\n",
            "\n",
            "🎯 Task: Get sales and create report\n",
            "\n",
            "📚 Found 1 similar tasks\n",
            "\n",
            "📋 Plan:\n",
            "Here's a step-by-step plan to get sales data and create a report, using the specified tools and format:\n",
            "\n",
            "**Plan:**\n",
            "\n",
            "**Step 1:** Retrieve the sales data from the specified source. [TOOL: fetch(source=\"sales_data\")]\n",
            "**Step 2:** Analyze the fetched sales data to identify key trends and metrics. [TOOL: analyze(data=\"sales_data\")]\n",
            "**Step 3:** Calculate total sales revenue from the analyzed data. [TOOL: calc(expr=\"total_revenue\")]\n",
            "**Step 4:** Generate a report summarizing the sales analysis and total revenue. [TOOL: report(msg=\"Sales analysis report: Total revenue = total_revenue\")]\n",
            "\n",
            "\n",
            "📚 Learned 2 tasks\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "N6AoyJyjalMc"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}